package cn.getech.data.development.sink.kudu

import java.io.{ByteArrayOutputStream, ObjectOutputStream}

import cn.getech.data.development.sink.jdbc.config.CustomTableConfig
import cn.getech.data.development.utils.JDBCSQLTypeUtils
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.{RichSinkFunction, SinkFunction}
import org.apache.flink.types.Row
import org.apache.kudu.client.{KuduClient, KuduSession, KuduTable, SessionConfiguration}
import org.slf4j.{Logger, LoggerFactory}

/**
 * kudu通用Sink
 * @param conf 配置参数
 */
class FlinkStreamKuduSink(conf: CustomTableConfig) extends RichSinkFunction[Row] {
  private val logger: Logger = LoggerFactory.getLogger(this.getClass)
  private var client: KuduClient = null
  private var table: KuduTable = null
  private var kuduSession: KuduSession = null
  private var out: ByteArrayOutputStream = null
  private var os: ObjectOutputStream = null

  override def open(parameters: Configuration): Unit = {
    val config = conf.getConfig()
    out = new ByteArrayOutputStream
    os = new ObjectOutputStream(out)
    client = new KuduClient.KuduClientBuilder(s"${config.url._2}").build
    if (client.tableExists(config.tablename._2)) {
      table = client.openTable(config.tablename._2)
    }
    kuduSession = client.newSession
    kuduSession.setFlushMode(SessionConfiguration.FlushMode.MANUAL_FLUSH)
    kuduSession.setMutationBufferSpace(conf.getBatchInterval)
  }

  override def invoke(value: Row, context: SinkFunction.Context): Unit = {
    val insert = table.newInsert
    val row = insert.getRow
    val utils = new JDBCSQLTypeUtils
    for (i <- 0 until conf.getParamSize) {
      utils.kudu_type2DataType(row, conf.getParams(i), value.getField(i).toString)
    }
    kuduSession.flush
    val response = kuduSession.apply(insert)
    if (response != null) logger.error(response.getRowError.toString)
  }

  override def close(): Unit = {
    try {
      kuduSession.close
      client.close()
      os.close()
      out.close
    } catch {
      case e: Exception =>
        logger.error(e.getMessage)
    }
  }
}
